#pragma once

// https://blog.csdn.net/fulianzhou/article/details/125186185

#include <cstdint>
#include <string>
#include <fstream>
#include <vector>
#include <memory>
#include <iostream>
#include <arpa/inet.h>
#include <map>
#include <array>
#include <cstring>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>

// pcap 文件头

#pragma pack(1)
struct PcapFileHeader
{
    uint32_t magic; // 0xa1b2c3d4
    uint16_t version_major;
    uint16_t version_minor;
    int32_t thiszone;  // gmt to local correction
    uint32_t sigfigs;  // accuracy of timestamps
    uint32_t snaplen;  // max length saved portion of each pkt
    uint32_t linktype; // data link type (LINKTYPE_*)
};

struct PcapTimeVal
{
    uint32_t tv_sec;  // seconds
    uint32_t tv_usec; // microseconds
};

// pcap 数据包头
struct PcapPacketHeader
{
    PcapTimeVal ts;  // time stamp
    uint32_t caplen; // length of portion present
    uint32_t len;    // length this packet (off wire)
};

// ethernet 数据帧头
struct EthernetFrameHeader
{
    uint8_t dst_mac[6];
    uint8_t src_mac[6];
    uint16_t type;
};

// ipv4 数据包头
struct Ipv4PacketHeader
{
#if __BYTE_ORDER == __LITTLE_ENDIAN
    uint8_t header_length : 4;
    uint8_t version : 4;
#elif __BYTE_ORDER == __BIG_ENDIAN
    uint8_t version : 4;
    uint8_t header_length : 4;
#endif
    uint8_t tos;
    uint16_t total_length;
    uint16_t id;
    union
    {
        struct
        {
#if __BYTE_ORDER == __LITTLE_ENDIAN
            uint16_t fragment_offset : 13;
            uint16_t flags : 3;
#elif __BYTE_ORDER == __BIG_ENDIAN
            uint16_t flags : 3;
            uint16_t fragment_offset : 13;
#endif
        };
        uint16_t all;
    } flags_fragment_offset;

    uint8_t ttl;
    uint8_t protocol;
    uint16_t checksum;
    uint32_t src_ip;
    uint32_t dst_ip;

    void ntoh()
    {
        total_length = ntohs(total_length);
        id = ntohs(id);
        flags_fragment_offset.all = ntohs(flags_fragment_offset.all);
        checksum = ntohs(checksum);
        src_ip = ntohl(src_ip);
        dst_ip = ntohl(dst_ip);
    }
    void hton()
    {
        total_length = htons(total_length);
        id = htons(id);
        flags_fragment_offset.all = htons(flags_fragment_offset.all);
        checksum = htons(checksum);
        src_ip = htonl(src_ip);
        dst_ip = htonl(dst_ip);
    }
};

// udp 数据包头
struct UdpPacketHeader
{
    uint16_t src_port;
    uint16_t dst_port;
    uint16_t length;
    uint16_t checksum;
    void hton()
    {
        src_port = htons(src_port);
        dst_port = htons(dst_port);
        length = htons(length);
        checksum = htons(checksum);
    }
    void ntoh()
    {
        src_port = ntohs(src_port);
        dst_port = ntohs(dst_port);
        length = ntohs(length);
        checksum = ntohs(checksum);
    }
};

// tcp 数据包头
struct TcpPacketHeader
{
    uint16_t src_port;
    uint16_t dst_port;
    uint32_t seq_num;
    uint32_t ack_num;
    uint8_t reserved : 4;
    uint8_t header_length : 4;
    uint8_t flags;
    uint16_t window_size;
    uint16_t checksum;
    uint16_t urgent_pointer;
    void ntoh()
    {
        src_port = ntohs(src_port);
        dst_port = ntohs(dst_port);
        seq_num = ntohl(seq_num);
        ack_num = ntohl(ack_num);
        window_size = ntohs(window_size);
        checksum = ntohs(checksum);
        urgent_pointer = ntohs(urgent_pointer);
    }
};

#pragma pack()

class PcapPaser
{
private:
    int m_fd;
    const std::string m_file_path;
    std::vector<uint8_t> m_content;
    struct stat m_file_info;

    void *m_file_content = nullptr;

    uint8_t *m_current = nullptr;
    uint8_t *m_end = nullptr;

    std::map<uint16_t, std::vector<std::vector<uint8_t>>> m_ip_fragment;

public:
    PcapPaser(const std::string &file_path) : m_file_path(file_path)
    {
    }
    ~PcapPaser()
    {
        ::munmap(m_file_content, m_file_info.st_size);
        ::close(m_fd);
    }

    bool open();
    bool isEOF() const
    {
        return m_current >= m_end;
    }

    size_t getFileSize() const{
        return m_file_info.st_size;
    }

    uint8_t *getIP(const std::string src,
                   const std::string dest,
                   int protocol);

    uint8_t *getTCP(const std::string src,
                    const std::string dest,
                    size_t* tcp_data_len);

    uint8_t *getUDP(const std::string src,
                    const std::string dest,
                    size_t* udp_data_len);
};
